//      ompr.h
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#ifndef ORTHOGONAL_MATCHING_PURSUIT_WITH_REPLACEMENT_H
#define ORTHOGONAL_MATCHING_PURSUIT_WITH_REPLACEMENT_H

#include "helper/stl_helper.h"
#include "helper/ublas_helper.h"
#include "helper/ml_data.h"

struct ele {
  double val;
  int idx;
  ele(double v=0, int i=0) : val(v), idx(i) {}

  struct greater_than {
    bool operator() (const ele& e1, const ele& e2) const { return e1.val > e2.val; }
  };

  struct less_than {
    bool operator() (const ele& e1, const ele& e2) const { return e1.val < e2.val; }
  };

  static std::vector<ele> prepare_sort(const vec_t& v, const vi_t& idx, bool absolute_value) {
    std::vector<ele> vec;
    for (int i = 0; i < sz(idx); i++)
      vec.pb(ele(absolute_value ? fabs(v(idx[i])) : v(idx[i]), idx[i]));
    return vec;
  }
};

template <typename D>
class OMPR {
public:
  typedef typename D::data_t data_t;
  typedef typename D::x_t x_t;

  boost::program_options::variables_map& vm_;
  int l_;
  int k_;
  double eta_;
  int max_iter_;

  sparse_vec_t x_;

  const sparse_vec_t& result() const { return x_; }
public:
  OMPR(boost::program_options::variables_map& vm) : vm_(vm), l_(1), k_(1), eta_(0.1), max_iter_(10)
  {
    if (vm.count("ompr_l")) l_ = VM(vm, "ompr_l", int);
    if (vm.count("ompr_k")) k_ = VM(vm, "ompr_k", int);
    if (vm.count("optim_eta")) eta_ = VM(vm, "optim_eta", double);
    if (vm.count("optim_iter")) max_iter_ = VM(vm, "optim_iter", int);
  }

  void solve(D& d)
  {
    int dim = d.dim();
    if (k_ > dim) FATAL_ERROR("Check k_=" << k_ << " dim=" << dim);
    data_t& A = d.data();
    vec_t& b = d.label();
    sparse_vec_t x(dim), x_tmp(dim);
    vi_t idx, non_idx;

    // normalize data
    for (int i = 0; i < dim; i++) {
      double norm_i = ublas::norm_2(col(A, i));
      col(A, i) /= norm_i;
    }

    // initialize
    for (int i = 0; i < k_; i++) idx.pb(i); 
    for (int i = k_; i < dim; i++) non_idx.pb(i);
    
    // loop
    double norm_r_before = ublas::norm_2(b);
    for (int iter = 0; iter < max_iter_; iter++) {
      vec_t residual = b;
      ublas::axpy_prod(A, -x, residual, false);

      // partial thresholding (adding l_ element from non-selected indices)
      // TODO: use nearest neighbor instead of inner product
      vec_t anti_grad = ublas::prod(ublas::trans(A), residual);
      vec_t z = x + eta_ * anti_grad;
      idx = partial_thresholding(z, idx, non_idx, k_, l_);

      // Solve model (restricted on nonzeros in idx)
      double norm_r_after;
      vec_t x_sub = solve_sub_model(A, b, idx, norm_r_after);

      // residual
      MSG_VERBOSE(vm_, "iter=" << iter << " norm(r)=" << norm_r_before);
      if (norm_r_after > norm_r_before) {
	MSG_VERBOSE(vm_, "maybe oscilating, terminate");
	break;
      }

      // copy x_sub to x and prepare non_idx for next iteration
      x = copy_from_sub(dim, x_sub, idx, non_idx);

      if (fabs(norm_r_before - norm_r_after) < 1e-6) {
	MSG("Converged after " << iter << " iteration.");
	break;
      }
      norm_r_before = norm_r_after;
    }
    x_ = x; // save result
  }

  static vi_t max_idx(const vec_t& z, const vi_t& idx, int k) 
  {
    std::vector<ele> absz = ele::prepare_sort(z, idx, true);    
    //for (int i = 0; i < sz(idx); i++) absz.pb(ele(fabs(z(idx[i])), idx[i]));
    std::sort(all_element(absz), ele::greater_than());

    vi_t midx;
    for (int i = 0; i < k && i < sz(absz); i++) midx.pb(absz[i].idx);
    std::sort(all_element(midx));
    return midx;
  }

  static vi_t partial_thresholding(const vec_t& z, const vi_t& idx, const vi_t& non_idx, int k, int l)
  {
    vi_t top = max_idx(z, non_idx, l), midx = idx;
    midx.insert(midx.end(), all_element(top));
    midx = max_idx(z, midx, k);
    return midx;
  }

  static sparse_vec_t copy_from_sub(int dim, const vec_t& x_sub, const vi_t& idx, vi_t& non_idx) {
    sparse_vec_t x(dim);
    std::vector<bool> marker(dim, true);
    for (int i = 0; i < sz(idx); i++) {
      x.pb(idx[i], x_sub(i));
      marker[idx[i]] = false;
    }
    non_idx.clear();
    for (int i = 0; i < dim; i++) 
      if (marker[i]) non_idx.pb(i);
    return x;
  }
  
  static vec_t solve_sub_model(const data_t& A, const vec_t& b, const vi_t& idx, double& norm_r) {
    data_t A_sub(A.size1(), sz(idx));
    for (int i = 0; i < sz(idx); i++) col(A_sub, i) = col(A, idx[i]);
    vec_t x_sub = ublas::prod(trans(A_sub), b);
    matrix_t A_sub_T_A_sub = ublas::prod(trans(A_sub), A_sub);
    int chol = cholesky_decompose(A_sub_T_A_sub);
    if (chol > 0) FATAL_ERROR("Cannot do Cholesky decomposition");
    cholesky_solve(A_sub_T_A_sub, x_sub, ublas::lower());

    vec_t r = b;
    ublas::axpy_prod(A_sub, -x_sub, r, false);
    norm_r = ublas::norm_2(r);
    return x_sub;
  }
};

#endif
